import struct
import socket
import time
from typing import Dict, List

BTN_IP_MAPPING = {
    'WS1-Riveting': ["192.168.31.241"], 
    'WS1-Welding': ["192.168.31.242"], 
    'WS1-Gluing': ["192.168.31.243"], 
    'WS2-L1': ["192.168.31.244"], 
    'WS2-L2': ["192.168.31.245"], 
    'WS2-Testing': ["192.168.31.246"],
}


class ModbusTCPClient:
    def __init__(self, host: str, port: int = 5003, unit_id: int = 1):
        self.host = host          # 按钮IP地址
        self.port = port          # 端口号，默认5003
        self.unit_id = unit_id    # 单元标识符，默认1
        self.sock = None          # socket连接对象
        self.transaction_id = 0   # 事务ID，用于请求响应匹配
        self.connected = False    # 连接状态标志
    
    def connect(self) -> bool:
        """连接到ModbusTCP按钮"""
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.settimeout(5.0)
            self.sock.connect((self.host, self.port))
            self.connected = True
            return True
        except Exception as e:
            print(f"{self.host}连接失败: {e}")
            return False
    
    def disconnect(self):
        """断开连接"""
        if self.sock:
            self.sock.close()
            self.connected = False
    
    def _build_mbap_header(self, length: int) -> bytes:
        """构建MBAP头部"""
        self.transaction_id = (self.transaction_id + 1) % 65536
        return struct.pack('>HHHB', 
                          self.transaction_id,  # 事务标识符
                          0,                    # 协议标识符
                          length,               # 长度
                          self.unit_id)         # 单元标识符
    
    def read_holding_registers(self, address: int, count: int = 1) -> List[int]:
        """读取保持寄存器（功能码03）"""
        if not self.connected:
            raise ConnectionError(f"未连接到按钮{self.host}")
        
        # 构建PDU
        pdu = struct.pack('>BHH', 0x03, address, count)
        # 构建完整请求
        mbap_header = self._build_mbap_header(len(pdu) + 1)
        request = mbap_header + pdu
        
        # 发送请求
        self.sock.send(request)
        
        # 接收响应
        response = self.sock.recv(1024)
        
        # 解析响应
        if len(response) < 9:
            raise ValueError(f"{self.host}响应数据长度不足")
        
        # 解析MBAP头部
        trans_id, proto_id, length, unit_id = struct.unpack('>HHHB', response[:7])
        
        # 解析PDU
        func_code = response[7]
        if func_code == 0x03:  # 正常响应
            byte_count = response[8]
            data = response[9:9+byte_count]
            # 解析寄存器值
            values = []
            for i in range(0, byte_count, 2):
                value = struct.unpack('>H', data[i:i+2])[0]
                values.append(value)
            return values
        else:
            raise ValueError(f"{self.host}读取寄存器失败，功能码: {func_code}")
    
    def write_single_register(self, address: int, value: int) -> bool:
        """写单个保持寄存器（功能码06）"""
        if not self.connected:
            raise ConnectionError(f"未连接到按钮{self.host}")
        
        # 构建PDU
        pdu = struct.pack('>BHH', 0x06, address, value)
        # 构建完整请求
        mbap_header = self._build_mbap_header(len(pdu) + 1)
        request = mbap_header + pdu
        
        # 发送请求
        self.sock.send(request)
        
        # 接收响应
        response = self.sock.recv(1024)
        
        # 解析响应
        if len(response) < 12:
            raise ValueError(f"{self.host}响应数据长度不足")
        
        # 检查响应是否正确
        func_code = response[7]
        if func_code == 0x06:
            return True
        else:
            raise ValueError(f"{self.host}写入寄存器失败，功能码: {func_code}")
    
    def write_multiple_registers(self, address: int, values: List[int]) -> bool:
        """写多个保持寄存器（功能码16）"""
        if not self.connected:
            raise ConnectionError(f"未连接到按钮{self.host}")
        
        count = len(values)
        byte_count = count * 2
        
        # 构建寄存器数据
        reg_data = b''
        for value in values:
            reg_data += struct.pack('>H', value)
        
        # 构建PDU
        pdu = struct.pack('>BHHB', 0x10, address, count, byte_count) + reg_data
        # 构建完整请求
        mbap_header = self._build_mbap_header(len(pdu) + 1)
        request = mbap_header + pdu
        
        # 发送请求
        self.sock.send(request)
        
        # 接收响应
        response = self.sock.recv(1024)
        
        # 解析响应
        if len(response) < 12:
            raise ValueError(f"{self.host}响应数据长度不足")
        
        func_code = response[7]
        if func_code == 0x10:
            return True
        else:
            raise ValueError(f"{self.host}写入寄存器失败，功能码: {func_code}")


def get_lineID(btn_name: str) -> str:
    if btn_name not in BTN_IP_MAPPING:
        raise ValueError(f"未找到名为{btn_name}的按钮")
    return "line" + str(list(BTN_IP_MAPPING.keys()).index(btn_name) + 1)  # 按钮ID从1开始，所以需要加1

# 获取对应按钮id的名称，一号产线id=1
def get_name(btn_id: int) -> str:
    return list(BTN_IP_MAPPING.keys())[btn_id-1]


def get_ip(btn_name: str) -> List[str]:
    """获取产线对应的IP地址列表"""
    if btn_name not in BTN_IP_MAPPING:
        raise ValueError(f"未找到名为{btn_name}的按钮对应的IP地址")
    return BTN_IP_MAPPING[btn_name]

def get_line_name_by_ip(ip: str) -> str:
    """通过IP地址反向查找产线名称"""
    for line_name, ip_list in BTN_IP_MAPPING.items():
        if ip in ip_list:
            return line_name
    raise ValueError(f"未找到IP地址{ip}对应的产线")

def get_all_ips() -> List[str]:
    """获取所有按钮的IP地址列表"""
    all_ips = []
    for ip_list in BTN_IP_MAPPING.values():
        all_ips.extend(ip_list)
    return all_ips


def get_regs_with_client(client: ModbusTCPClient) -> Dict:
    """使用现有客户端连接获取呼叫盒的状态信息"""
    try:
        # 读取所有寄存器（0-29）
        regs = client.read_holding_registers(0, 30)
        
        return {
            "dev_id": regs[0],  # 呼叫盒编号
            "model": "ModbusTCP",
            "regs": {i: regs[i] for i in range(30)}
        }
    except Exception as e:
        raise ConnectionError(f"获取呼叫盒状态失败: {str(e)}")

def set_regs_with_client(client: ModbusTCPClient, reg_values: Dict[int, int]) -> Dict:
    """使用现有客户端连接修改呼叫盒的寄存器值"""
    for reg_id, value in reg_values.items():
        if not (10 <= reg_id <= 29):
            raise ValueError(f"寄存器ID {reg_id} 超出有效范围[10,29]")
        if not (0 <= value <= 29):
            raise ValueError(f"寄存器值 {value} 超出有效范围[0,29]")
    
    try:
        # 写入寄存器
        for reg_id, value in reg_values.items():
            client.write_single_register(reg_id, value)
        
        return {"Code": 0, "Msg": "Success"}
    except Exception as e:
        raise ConnectionError(f"修改呼叫盒寄存器失败: {str(e)}")

def blink_led_with_client(client: ModbusTCPClient, led_reg: int = 12, duration: float = 1.5) -> None:
    """使用现有客户端连接控制LED闪烁"""
    try:  
        # 设置LED为快闪（值为2）
        set_regs_with_client(client, {led_reg: 1})
        
        time.sleep(duration)
        
        # 设置LED为关闭（值为0）
        set_regs_with_client(client, {led_reg: 0})
        
    except Exception as e:
        print(f"LED{led_id}闪烁过程中出错: {str(e)}")